import numpy as np


def bjk(B, j, k):
    return B[j][k]


def fw(A, B, Pi, O, T, N):
    # size of a is len of N * T
    alpha = np.zeros([N, T])

    for t in range(T):
        for j in range(N):
            # if t = 0, then initial the value of a
            if t == 0:
                alpha[j][t] = Pi[j] * bjk(B, j, O[t])
            else:
                for i in range(N):
                    #a_t+1(j) = sum of a_t(i)* aij * B_j(O_t+1) 
                    alpha[j][t] += alpha[i][t - 1] * A[i][j] * bjk(B, j, O[t])
    # p(O|λ) = sum(alpha)        
    P_O_lambda = alpha.sum(axis=0)[T-1]
    return P_O_lambda


def bw(A, B, Pi, O, T, N):
    beta = np.zeros([N, T])
    for t in reversed(range(T)):
        for i in range(N):
            #set initial value = 1
            if t == T-1:
                beta[i][t] = 1
            else:
                for j in range(N):
                    #update b_t(i) by sum of B_j(o_t+1) * b_t+1(j) * aij
                    beta[i][t] += bjk(B, j, O[t+1]) * beta[j][t+1] * A[i][j]
    
    res = 0
    for i in range(N):        
        res += Pi[i] * bjk(B, i, O[0]) * beta[i][0]
    
    return res



if __name__ == '__main__':
    # given lambda
    A = np.array([[0.5, 0.2, 0.3], [0.3, 0.5, 0.2], [0.2, 0.3, 0.5]])
    B = np.array([[0.5, 0.5], [0.4, 0.6], [0.7, 0.3]])
    Pi = np.transpose(np.array([0.2, 0.4, 0.4]))
    # give O
    O = [0, 1, 0]
    T = 3
    N = len(A)
    # Forward algorithm
    fw_res = fw(A, B, Pi, O, T, N)
    # backward algorithm
    bw_res = bw(A, B, Pi, O, T, N)


